package com.study.feature.transform

import org.apache.spark.ml.feature.{RegexTokenizer, Tokenizer}
import org.apache.spark.sql.SparkSession

/**
 * 特征转换-分词
 *
 * @author stephen
 * @date 2019-08-28 11:27
 */
object TokenizerDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")

    import org.apache.spark.sql.functions._

    val sentenceDataFrame = spark.createDataFrame(Seq(
      (0, "Hi I heard about Spark"),
      (1, "I wish Java could use case classes"),
      (2, "Logistic,regression,models,are,neat")
    )).toDF("id", "sentence")

    // 默认配置的分词器
    val tokenizer = new Tokenizer().setInputCol("sentence").setOutputCol("words")
    // 基于正则表达式的分词器
    val regexTokenizer = new RegexTokenizer()
      .setInputCol("sentence")
      .setOutputCol("words")
      .setPattern("\\W") // 也可以换成 .setPattern("\w+").setGaps(false)

    val countTokens = udf { (words: Seq[String]) => words.length }

    // 默认分词器进行分词
    val tokenized = tokenizer.transform(sentenceDataFrame)
    tokenized.select("sentence", "words")
      .withColumn("tokens", countTokens(col("words"))).show(false)

    // 基于正则表达式的分词器分词
    val regexTokenized = regexTokenizer.transform(sentenceDataFrame)
    regexTokenized.select("sentence", "words")
      .withColumn("tokens", countTokens(col("words"))).show(false)

    spark.stop()
  }
}
